from flask import Flask, request, jsonify, send_file
import pandas as pd
import os
from werkzeug.utils import secure_filename
import json

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['ALLOWED_EXTENSIONS'] = {'xlsx', 'xls', 'csv'}

# 确保上传目录存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in app.config['ALLOWED_EXTENSIONS']

@app.route('/')
def index():
    return send_file('index.html')

@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'error': '没有文件被上传'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': '没有选择文件'}), 400
    
    if file and allowed_file(file.filename):
        filename = secure_filename(file.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        
        try:
            # 读取文件
            if filename.endswith('.csv'):
                df = pd.read_csv(filepath)
            else:
                df = pd.read_excel(filepath)
            
            # 获取列信息
            columns = df.columns.tolist()
            data_preview = df.head(10).to_dict('records')
            
            # 为每列生成唯一值（用于下拉选项）
            column_options = {}
            for col in columns:
                unique_values = df[col].dropna().unique().tolist()
                column_options[col] = unique_values
            
            return jsonify({
                'columns': columns,
                'data_preview': data_preview,
                'column_options': column_options,
                'filename': filename
            })
            
        except Exception as e:
            return jsonify({'error': f'文件解析错误: {str(e)}'}), 500
    
    return jsonify({'error': '不支持的文件格式'}), 400

@app.route('/export', methods=['POST'])
def export_file():
    data = request.json
    filename = data.get('filename')
    filters = data.get('filters', {})
    
    if not filename:
        return jsonify({'error': '缺少文件名'}), 400
    
    filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    
    try:
        # 重新读取原始文件
        if filename.endswith('.csv'):
            df = pd.read_csv(filepath)
        else:
            df = pd.read_excel(filepath)
        
        # 应用筛选条件
        filtered_df = df.copy()
        for column, selected_values in filters.items():
            if selected_values:
                filtered_df = filtered_df[filtered_df[column].isin(selected_values)]
        
        # 保存筛选后的文件
        output_filename = f"filtered_{filename}"
        output_path = os.path.join(app.config['UPLOAD_FOLDER'], output_filename)
        
        if filename.endswith('.csv'):
            filtered_df.to_csv(output_path, index=False)
        else:
            filtered_df.to_excel(output_path, index=False)
        
        return send_file(output_path, as_attachment=True, download_name=output_filename)
        
    except Exception as e:
        return jsonify({'error': f'文件导出错误: {str(e)}'}), 500

if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)
